import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from tqdm import tqdm
import argparse
import base64
import json
import sys
import os
import os.path as osp
import io

import imgviz
import PIL.Image
import yaml
import pandas as pd

# image to array
def img_b64_to_arr(img_b64):
    f = io.BytesIO()
    f.write(base64.b64decode(img_b64))
    img_arr = np.array(PIL.Image.open(f))
    return img_arr

# Convert labels shape to array
def shapes_to_label(img_shape, shapes, label_name_to_value, type='class'):    
    assert type in ['class', 'instance']
    mask_shape_array = np.zeros(img_shape[:2], dtype=np.int32)
    
    if type == 'instance':
        ins = np.zeros(img_shape[:2], dtype=np.int32)
        instance_names = ['_background_']
    
    for shape in shapes:
        
        points = shape['points'] # mask points
        label = shape['label'] # mask label name -> String
        shape_type = shape.get('shape_type', None) # polygon or something else
        
        if type == 'class':
            cls_name = label # mask label name -> String
            
        elif type == 'instance':
            cls_name = label.split('-')[0]
            if label not in instance_names:
                instance_names.append(label)
            ins_id = instance_names.index(label)
        
        cls_name = cls_name.strip()
        class_id = label_name_to_value[cls_name] # mask label id in 0, 1, 2...
        mask = shape_to_mask(img_shape[:2], points, shape_type) # mask shape (H * W)
        mask_shape_array[mask] = class_id # set the area mask not zero to mask_id
        
        if type == 'instance':
            ins[mask] = ins_id

    if type == 'instance':
        return mask_shape_array, ins
    
    return mask_shape_array

# Convert shape to mask
def shape_to_mask(img_shape, points, shape_type=None, line_width=10, point_size=5):
    
    mask = np.zeros(img_shape[:2], dtype=np.uint8) # generate numpy (H * W)

    mask = PIL.Image.fromarray(mask) # convert nupmy to PIL

    draw = PIL.ImageDraw.Draw(mask) 
    
    xy = [tuple(point) for point in points] # generate (x,y)
    
    if shape_type == 'circle':
        assert len(xy) == 2, 'Shape of shape_type=circle must have 2 points'
        (cx, cy), (px, py) = xy
        d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
        draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1) 
    elif shape_type == 'rectangle':
        assert len(xy) == 2, 'Shape of shape_type=rectangle must have 2 points'
        draw.rectangle(xy, outline=1, fill=1)
    elif shape_type == 'line':
        assert len(xy) == 2, 'Shape of shape_type=line must have 2 points'
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == 'linestrip':
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == 'point':
        assert len(xy) == 1, 'Shape of shape_type=point must have 1 points'
        cx, cy = xy[0]
        r = point_size
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
    
    else:
        assert len(xy) > 2, 'Polygon must have points more than 2'
        
        draw.polygon(xy=xy, outline=1, fill=1)
    
    mask = np.array(mask, dtype=bool)
  
    return mask


# Save mask to image
def save_mask_to_image(filename, mask_with_class_id):
    
    if osp.splitext(filename)[1] != '.png':
        filename += '.png'

    # Assume label ranses [-1, 254] for int32, and [0, 255] for uint8 as VOC.
    if mask_with_class_id.min() >= -1 and mask_with_class_id.max() < 255:
        mask_with_class_id_pil = PIL.Image.fromarray(mask_with_class_id.astype(np.uint8), mode='P')
        
        colormap = imgviz.label_colormap(n_label=128) # generate color map
        
        mask_with_class_id_pil.putpalette(colormap.flatten()) # save mask to image conposed with diffierent color
        mask_with_class_id_pil.save(filename)
        
    else:
        raise ValueError(
            '[%s] Cannot save the pixel-wise class label as PNG. '
            'Please consider using the .npy format.' % filename
        )

def json_to_image(img_path, json_path, toImage=False, visual=False, mask_label='box'):
    """ Covert json labels to segementation masks
    Args:
        img_path: input img path
        json_path: json file full name
        toImage:
        visual: whether show the mask on origin image
        mask_label: label name for the mask
    Returns:
        save masks to given output folder
    Raises:
        IOError: An error occurred accessing ablumentation object.
    """
    
    with open(json_path, 'r+', encoding="utf-8") as f:
        """ Load json file and update dictionary order """
        data = json.load(f)
        data['shapes'] = sorted(data['shapes'], key = lambda x: x['label'], reverse=False)
        imageData = data.get('imageData')
        # data = json.load(open(os.path.join(path_in, json_file)), encoding="utf-8-sig")
        
        """ Convert imageData to numpy array """
        if not imageData:
            if Path(data['imagePath']).stem == Path(img_path).stem:
                img = np.asarray(PIL.Image.open(img_path))
                # with open(img_path, 'rb') as f:
                #     imageData = f.read()
                #     imageData = base64.b64encode(imageData).decode('utf-8')
                #     img = img_b64_to_arr(imageData)
            else:
                raise ValueError('Can not find image file for this mask file.')
        else:
            img = img_b64_to_arr(imageData)
        # print(f"01 The dimension of origin image is {img.shape}")

        """ Generate labels in the type of label encoder """
        mask_to_label_encode = {'_background_': 0, mask_label: 1} # label encode each mask class
        # mask_to_label_encode = {'_background_': 0, '01_bonding': 1, 
        #                         '02_spot': 2, '03_spatter': 3, '04_slag': 4}
        
        
        #         for shape in sorted(data['shapes'], key = lambda x: x['label'], reverse=False): # data['shapes'] -> list, sort by label class
        #             mask_label_name = shape['label'] # Only one shape each round

        #             if mask_label_name in mask_to_label_encode:
        #                 mask_label_value = mask_to_label_encode[mask_label_name] # mask_label_name means 'Joint', 'Line', or something else
        #             else:
        #                 mask_label_value = len(mask_to_label_encode) # If mask not include in list, add new mask label value
        #                 mask_to_label_encode[mask_label_name] = mask_label_value

        """ Generate labels in the type of classID """        
        mask_with_class_id = shapes_to_label(img.shape, data['shapes'], mask_to_label_encode) # mask's value are class id
        
        return mask_with_class_id


def json_to_image_one_by_one(img_shape, json_path, toImage=False, visual=False, mask_label='box'):
    """ Covert json labels to segementation masks
    Args:
        img_shape: [h, w, c]
        json_path: json file full name
        toImage:
        visual: whether show the mask on origin image
        mask_label: label name for the mask
    Returns:
        save masks to given output folder
    Raises:
        IOError: An error occurred accessing ablumentation object.
    """
    
    with open(json_path, 'r+', encoding="utf-8") as f:
        """ Load json file and update dictionary order """
        data = json.load(f)

    data['shapes'] = sorted(data['shapes'], key = lambda x: x['label'], reverse=False)
    # imageData = data.get('imageData')
    # # data = json.load(open(os.path.join(path_in, json_file)), encoding="utf-8-sig")
    
    # """ Convert imageData to numpy array """
    # if not imageData:
    #     if Path(data['imagePath']).stem:
    #         img = np.asarray(Image.open(img_path))
    #         # with open(img_path, 'rb') as f:
    #         #     imageData = f.read()
    #         #     imageData = base64.b64encode(imageData).decode('utf-8')
    #         #     img = img_b64_to_arr(imageData)
    #     else:
    #         raise ValueError('Can not find image file for this mask file.')
    # else:
    #     img = img_b64_to_arr(imageData)
    # print(f"01 The dimension of origin image is {img.shape}")

    """ Generate labels in the type of label encoder """
    # label encode each mask class
    mask_to_label_encode = {'_background_': 0}
    for i in range(len(data['shapes'])):
        shape_slice = data['shapes'][i]
        label_slice = shape_slice['label']
        if label_slice not in mask_to_label_encode.keys():
            mask_to_label_encode[label_slice] = 1
    # mask_to_label_encode = {'_background_': 0, mask_label: 1} 
    # mask_to_label_encode = {'_background_': 0, '01_bonding': 1, 
    #                         '02_spot': 2, '03_spatter': 3, '04_slag': 4}
    
    
    #         for shape in sorted(data['shapes'], key = lambda x: x['label'], reverse=False): # data['shapes'] -> list, sort by label class
    #             mask_label_name = shape['label'] # Only one shape each round

    #             if mask_label_name in mask_to_label_encode:
    #                 mask_label_value = mask_to_label_encode[mask_label_name] # mask_label_name means 'Joint', 'Line', or something else
    #             else:
    #                 mask_label_value = len(mask_to_label_encode) # If mask not include in list, add new mask label value
    #                 mask_to_label_encode[mask_label_name] = mask_label_value

    """ Generate labels in the type of classID """
    for s in data['shapes']:
        label = s['label']  
        mask_with_class_id = shapes_to_label(img_shape, [s], mask_to_label_encode) # mask's value are class id
    
        yield label, mask_with_class_id
